# coding: utf-8
import torch

@torch.jit.script
def jitdecode(loc, priors, variances):
    # type: (Tensor, Tensor, List[float]) ->Tensor
    """Decode locations from predictions using priors to undo
    the encoding we did for offset regression at train time.
    Args:
        loc (tensor): location predictions for loc layers,
            Shape: [num_priors,4]
        priors (tensor): Prior boxes in center-offset form.
            Shape: [num_priors,4].
        variances: (list[float]) Variances of priorboxes
    Return:
        decoded bounding box predictions
    """
    xy = priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:]
    wh = priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])

    lt = xy - wh / 2
    rb = lt + wh

    return torch.cat([lt, rb], 1)

@torch.jit.script
def jitcenter_size(boxes):
    # type:(Tensor) -> Tensor
    """ Convert prior_boxes to (cx, cy, w, h)
    representation for comparison to center-size form ground truth data.
    Args:
        boxes: (tensor) point_form boxes
    Return:
        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
    """
    cxy = (boxes[:, 2:] + boxes[:, :2]) / 2
    wh = boxes[:, 2:] - boxes[:, :2]
    return torch.cat([cxy, wh], 1)


class DetectionOutput(torch.jit.ScriptModule):
    __constants__ = ["num_classes"
        , "topk"
        , "conf_thresh"
        , "nms_thresh"
        , "variance"
        , "object_score"
        , "device"
        , "bkg_label", "background_label"]

    def __init__(self, num_classes, bkg_label, topk, conf_thresh, nms_thresh, variance, object_score=0.01):
        super(DetectionOutput, self).__init__()
        self.num_classes = num_classes
        self.background_label = bkg_label
        self.topk = topk
        # Parameters used in nms.
        self.nms_thresh = nms_thresh
        self.conf_thresh = conf_thresh
        self.variance = variance
        self.object_score = object_score

    @torch.jit.script_method
    def forward(self, arm_loc, arm_conf, odm_loc, odm_conf, priors):
        # 1, arm conf
        arm_object_conf = arm_conf[:, :, 1] > self.object_score
        arm_object_conf = arm_object_conf.unsqueeze(2).expand_as(odm_conf).float()
        odm_conf = odm_conf * arm_object_conf
        num = odm_loc.size(0)
        num_priors = priors.size(0)
        conf_preds = odm_conf.view(num, num_priors, self.num_classes).transpose(2, 1)
        output = []
        for i in range(num):
            num_out_loc = []
            num_out_class = []
            # first arm loc
            default = jitdecode(arm_loc[i], priors, self.variance)
            default = jitcenter_size(default)
            decoded_boxes = jitdecode(odm_loc[i], default, self.variance)
            # For each class, perform nms
            conf_scores = conf_preds[i].clone()
            for cl in range(self.num_classes - 1):
                c_mask = conf_scores[cl + 1].gt(self.conf_thresh)
                scores = conf_scores[cl + 1][c_mask]
                if scores.numel() > 0:
                    l_mask = c_mask.unsqueeze(1).expand_as(decoded_boxes)
                    boxes = decoded_boxes[l_mask].view(-1, 4)
                    # idx of highest scoring and non-overlapping boxes per class
                    ids = self.jitnms(boxes, scores)
                    count = ids.size(0)
                    a = scores[ids[:count]].unsqueeze(1)
                    # print(a)
                    b = boxes[ids[:count]]
                    c = torch.cat([a, b], 1)
                    num_out_loc += [c]
                    num_out_class += [torch.zeros(1, device=c.device).float() + cl + 1]

            num_out = []
            for j in range(len(num_out_loc)):
                a = num_out_loc[j]
                b = num_out_class[j]
                for k in range(a.size(0)):
                    c = torch.cat([b, a[k]], 0)
                    num_out += [c]
            num_out = torch.cat(num_out, 0)
            num_out = num_out.contiguous().view(1, -1, 6)
            output.append(num_out)

        # f = torch.cat(output, 0)
        return torch.cat(output, 0)

    @torch.jit.script_method
    def jitnms(self, boxes, scores):
        # (Tensor, Tensor) -> Tuple[Tensor, int]
        # print("intot nms")
        # keep = torch.zeros(scores.size(0)).long()
        x1 = boxes[:, 0]
        y1 = boxes[:, 1]
        x2 = boxes[:, 2]
        y2 = boxes[:, 3]
        area = torch.mul(x2 - x1, y2 - y1)
        v, idx = scores.sort(0)  # sort in ascending order
        idx = idx[-self.topk:]  # indices of the top-k largest vals

        xx1 = scores[0].clone()
        yy1 = scores[0].clone()
        xx2 = scores[0].clone()
        yy2 = scores[0].clone()
        w = scores[0].clone()
        h = scores[0].clone()

        # count = torch.zeros(1).long()
        # count = 0
        keep = []
        while idx.numel() > 1:
            i = idx[-1]  # index of current largest val
            # keep.append(i)
            # keep[count] = i
            # count += 1
            keep += [(torch.zeros(1) + i).long()]

            idx = idx[:-1]  # remove kept element from view
            # load bboxes of next highest vals
            torch.index_select(x1, 0, idx, out=xx1)
            torch.index_select(y1, 0, idx, out=yy1)
            torch.index_select(x2, 0, idx, out=xx2)
            torch.index_select(y2, 0, idx, out=yy2)
            # store element-wise max with next highest score
            xx1 = torch.clamp(xx1, min=x1[i].item())
            yy1 = torch.clamp(yy1, min=y1[i].item())
            xx2 = torch.clamp(xx2, max=x2[i].item())
            yy2 = torch.clamp(yy2, max=y2[i].item())
            w.resize_as_(xx2)
            h.resize_as_(yy2)
            w = xx2 - xx1
            h = yy2 - yy1
            # check sizes of xx1 and xx2.. after each iteration
            w = torch.clamp(w, min=0.0)
            h = torch.clamp(h, min=0.0)
            inter = w * h
            # IoU = i / (area(a) + area(b) - i)
            rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)
            union = (rem_areas - inter) + area[i]
            IoU = inter / union  # store result in iou
            # keep only elements with an IoU <= overlap
            le = IoU.le(self.nms_thresh)
            # select not support in jit script
            # idx = idx[le]
            # print(idx)
            tmp = []
            for j in range(le.size(0)):
                if bool(le[j] == 1):
                    tmp += [idx[j] + torch.zeros(1).long()]
            idx = torch.cat(tmp, 0)


        res = torch.cat(keep, 0)
        return res

